from .base_strategy import BaseStrategy
import backtrader as bt
class ThreeBullishStrategy(BaseStrategy):
    params = (
       ('ema_short', 20),
        ('ema_long', 120),
        ('consecutive_candles', 3),
    )
 
    def __init__(self):
        super().__init__()

        self.init_indicators()
        # 用于跟踪订单
        self.order = None
    
    def init_indicators(self):
        """在此定义指标 (EMA/MACD/RSI等)"""
        # 计算 EMA 指标
        self.ema_short = bt.indicators.EMA(self.data.close, period=self.p.ema_short)
        self.ema_long = bt.indicators.EMA(self.data.close, period=self.p.ema_long)
 
        # 用于记录连续阳线或阴线的计数
        self.consecutive_bullish = 0
        self.consecutive_bearish = 0


    def notify_order(self, order):
        super().notify_order(order)
 
    def next(self):
        super().next()
 
        # 判断趋势
        above_ema_long = self.data.close[0] > self.ema_long[0] # 多头趋势
        below_ema_long = self.data.close[0] < self.ema_long[0] # 空头趋势
 
        # 检查是否满足开仓条件
        if not self.position:
            if above_ema_long:
                self.check_bullish_condition()
            elif below_ema_long:
                self.check_bearish_condition()
        else:
            self.check_exit_condition() # 检查是否平仓
 
    def check_bullish_condition(self):
        # 检查是否为阳线，且没有或仅有很短的上下影线
        body = self.data.close[0] - self.data.open[0]
        upper_shadow = self.data.high[0] - max(self.data.open[0], self.data.close[0])
        lower_shadow = min(self.data.open[0], self.data.close[0]) - self.data.low[0]
 
        if body > 0 : # and upper_shadow < 0.3 * body and lower_shadow < 0.3 * body:
            self.consecutive_bullish += 1
            if self.consecutive_bullish >= self.p.consecutive_candles and self.data.close[0] > self.ema_short[0]:
                self.buy()
        else:
            self.consecutive_bullish = 0
 
    def check_bearish_condition(self):
        # 检查是否为阴线，且没有或仅有很短的上下影线
        body = self.data.open[0] - self.data.close[0]
        upper_shadow = self.data.high[0] - max(self.data.open[0], self.data.close[0])
        lower_shadow = min(self.data.open[0], self.data.close[0]) - self.data.low[0]
 
        if body > 0 :#and upper_shadow < 0.3 * body and lower_shadow < 0.3 * body:
            self.consecutive_bearish += 1
            if self.consecutive_bearish >= self.p.consecutive_candles and self.data.close[0] < self.ema_short[0]:
                self.sell()
        else:
            self.consecutive_bearish = 0
 
    def check_exit_condition(self):
        # 计算最近3根K线的最低价
        if len(self.data) >= 3:
            lowest_low = min(self.data.low[-3], self.data.low[-2], self.data.low[-1])
            lowest_high = max(self.data.high[-3], self.data.high[-2], self.data.high[-1])
            if self.position.size > 0:  # 多头仓位
                if self.data.close[0] < lowest_low*1.02:
                    self.logger.info(f'多头平仓 - 日期: {self.data.datetime.datetime(0)}, 最低价: {lowest_low:.2f}, 收盘价: {self.data.close[0]:.2f}')
                    self.close()
            elif self.position.size < 0:  # 空头仓位

                if self.data.close[0] > lowest_high*1.02:
                    self.logger.info(f'空头平仓 - 日期: {self.data.datetime.datetime(0)}, 最高价: {lowest_high:.2f}, 收盘价: {self.data.close[0]:.2f}')
                    self.close()